import os

from PyQt5.QtCore import pyqtSlot, QTimer, Qt
from PyQt5.QtGui import QImage, QPixmap, QPalette

from savePersonInfo import savePersonInfo
from ui.Ui_MianWindow import Ui_MainWindow
from openni import openni2
from tools import *

import numpy as np
import cv2


class mainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        QtWidgets.QMainWindow.__init__(self)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.ui.stopButton.setEnabled(False)
        self.ui.exportButton.setEnabled(False)

        self.rgb_img = None
        self.depth_img = None
        self.depth_stream = None
        self.color_stream = None
        self.dev = None
        self.tempType = None
        self.data_count = 10

        self.left_temp_rgb_list = []
        self.left_temp_depth_list = []
        self.right_temp_rgb_list = []
        self.right_temp_depth_list = []

        self.timer_camera = QTimer()
        #
        self.ui.leftDataButton.clicked.connect(lambda: self.collectDataClicked(tempType.LEFT_TEMP))
        self.ui.rightDataButton.clicked.connect(lambda: self.collectDataClicked(tempType.RIGHT_TEMP))
        self.timer_camera.timeout.connect(self.openFrame)
        self.ui.stopButton.clicked.connect(self.stopCollectData)
        self.ui.exportButton.clicked.connect(self.savePersonInfo)

    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
        try:
            if openni2.is_initialized():
                self.timer_camera.stop()
                self.depth_stream.stop()
                self.color_stream.stop()
                self.dev.close()
                openni2.unload()
        except Exception as err:
            print(err)
        # QtWidgets.QMainWindow.closeEvent(self,a0)

    @pyqtSlot()
    def savePersonInfo(self):
        dialog = savePersonInfo()
        if dialog.exec():
            # for i in len(self.left_temp_rgb_list)
            folder_path = 'face_data/{}_{}_{}'.format(dialog.name, dialog.height, dialog.weight)
            if not os.path.exists(folder_path):
                os.makedirs(folder_path)
            for i in range(0, len(self.left_temp_rgb_list)):
                # print(self.left_temp_rgb_list[i])
                cv2.imencode('.jpg',self.left_temp_rgb_list[i])[1].tofile('{}/{}_left_rgb.jpg'.format(folder_path, str(i)))
            for i in range(0, len(self.left_temp_depth_list)):
                cv2.imencode('.jpg',self.left_temp_depth_list[i])[1].tofile('{}/{}_left_depth.jpg'.format(folder_path, str(i)))
            for i in range(0, len(self.right_temp_rgb_list)):
                cv2.imencode('.jpg',self.right_temp_rgb_list[i])[1].tofile('{}/{}_right_rgb.jpg'.format(folder_path, str(i)))
            for i in range(0, len(self.right_temp_depth_list)):
                cv2.imencode('.jpg',self.right_temp_depth_list[i])[1].tofile('{}/{}_right_depth.jpg'.format(folder_path, str(i)))
            dialog.close()

            self.clearData(tempType.LEFT_TEMP)
            self.clearData(tempType.RIGHT_TEMP)
            alert(msg="\n保存成功，数据保存在目录：\n{}\n".format(os.path.abspath(folder_path)), title="消息",
                  icon=QtWidgets.QMessageBox.Information)

    def openFrame(self):
        # 显示深度图
        frame = self.depth_stream.read_frame()
        dframe_data = np.array(frame.get_buffer_as_triplet()).reshape([480, 640, 2])
        dframe1 = np.asarray(dframe_data[:, :, 0], dtype='float32')
        dframe2 = np.asarray(dframe_data[:, :, 1], dtype='float32')
        dpt = dframe1 + dframe2 * 255
        dpt = np.flipud(dpt)
        vertices = convertDepthToPoint(dpt)

        img_depth = cv2.cvtColor(dpt, cv2.COLOR_BGR2RGB)
        self.depth_img = QImage(img_depth.data, img_depth.shape[1], img_depth.shape[0], QtGui.QImage.Format_RGB888)
        self.ui.label_4.setPixmap(QPixmap.fromImage(self.depth_img))
        self.ui.label_4.setAlignment(QtCore.Qt.AlignHCenter | QtCore.Qt.AlignVCenter)

        # 提取RGB图像
        cframe = self.color_stream.read_frame()
        cframe_data = np.array(cframe.get_buffer_as_triplet()).reshape(
            [intrinsic_params.height, intrinsic_params.width, 3])
        R = cframe_data[:, :, 0]
        G = cframe_data[:, :, 1]
        B = cframe_data[:, :, 2]
        cframe_data = np.transpose(np.array([B, G, R]), [1, 2, 0])
        image_rows, image_cols, _ = cframe_data.shape
        face_img_rgb = cv2.cvtColor(cframe_data, cv2.COLOR_BGR2RGB)
        # 画出标记点
        results = face_mesh.process(face_img_rgb)
        if not results.multi_face_landmarks:
            return
        face_landmarks = results.multi_face_landmarks[0]
        landmarks = []
        landmark_pxs = []
        temporal_indexs = left_temporal_indexs
        temporal_drawlist = left_temporal_drawlist
        label = self.ui.label
        msg_txt = '左颞部'

        if self.tempType is tempType.RIGHT_TEMP:
            temporal_indexs = right_temporal_indexs
            temporal_drawlist = right_temporal_drawlist
            label = self.ui.label_2
            msg_txt = '右颞部'

        mp_drawing.draw_landmarks(face_img_rgb, face_landmarks, temporal_drawlist, drawing_spec, drawing_spec1)
        try:
            for i in range(0, len(face_landmarks.landmark)):
                landmark_px = mp_drawing._normalized_to_pixel_coordinates(face_landmarks.landmark[i].x,
                                                                          face_landmarks.landmark[i].y,
                                                                          image_cols, image_rows)
                # 记得要翻转y轴坐标
                landmark = [landmark_px[0], intrinsic_params.height - landmark_px[1]]
                landmark_pxs.append(landmark)
        except:
            return

        self.rgb_img = QImage(face_img_rgb.data, face_img_rgb.shape[1], face_img_rgb.shape[0],
                              QtGui.QImage.Format_RGB888)
        self.ui.label_3.setPixmap(QPixmap.fromImage(self.rgb_img))
        self.ui.label_3.setAlignment(QtCore.Qt.AlignHCenter | QtCore.Qt.AlignVCenter)

        if len(landmark_pxs) == 0:
            return
        for index in temporal_indexs:
            x, y = landmark_pxs[index]
            landmarks.append(vertices[y * intrinsic_params.width + x])
        landmarks = np.array(landmarks)
        if np.any(landmarks[:, 2] == 0.0):
            print("has 0 value")
            return

        count = 0
        if self.tempType is tempType.LEFT_TEMP:
            # cv2.imshow("rgb",cframe_data)
            self.left_temp_rgb_list.append(cframe_data)
            self.left_temp_depth_list.append(dpt)
            count = len(self.left_temp_depth_list)
        else:
            self.right_temp_rgb_list.append(cframe_data)
            self.right_temp_depth_list.append(dpt)
            count = len(self.right_temp_rgb_list)

        label.setText("{}：{}/10".format(msg_txt, str(count)))
        print(count)
        if 0 < count < self.data_count:
            label.setStyleSheet(start_collect_styleSheet)
        if count == self.data_count:
            label.setStyleSheet(end_collect_styleSheet)
            label.setText("完成{}采集，共{}张".format(msg_txt, str(count)))
            self.stopCollectData()

    def clearData(self, tempType):
        if tempType is tempType.RIGHT_TEMP:
            self.right_temp_rgb_list = []
            self.right_temp_depth_list = []
            self.ui.label_2.setStyleSheet(none_collect_styleSheet)
            self.ui.label_2.setText("右颞部：0/10")
        else:
            self.left_temp_rgb_list = []
            self.left_temp_depth_list = []
            self.ui.label.setStyleSheet(none_collect_styleSheet)
            self.ui.label.setText("左颞部：0/10")
        self.ui.exportButton.setEnabled(False)

    @pyqtSlot()
    def collectDataClicked(self, tempType: tempType):

        self.ui.rightDataButton.setEnabled(False)
        self.ui.leftDataButton.setEnabled(False)
        self.ui.stopButton.setEnabled(True)
        self.tempType = tempType
        self.clearData(tempType)
        try:
            openni2.initialize("thirdParty/obc_x64/Redist")
            # self.dev = openni2.Device.open_file('test.oni'.encode('utf-8'))
            self.dev = openni2.Device.open_any()
            print(self.dev.get_device_info())
            self.depth_stream = self.dev.create_depth_stream()
            self.color_stream = self.dev.create_color_stream()

            self.depth_stream.start()
            self.color_stream.start()
            self.timer_camera.start(100)
        except Exception as err:
            print(err)
            alert(str(err))
            self.ui.rightDataButton.setEnabled(True)
            self.ui.leftDataButton.setEnabled(True)
            self.ui.stopButton.setEnabled(False)

    @pyqtSlot()
    def stopCollectData(self):
        try:
            self.depth_stream.stop()
            self.color_stream.stop()
            self.dev.close()
            openni2.unload()
        except Exception as err:
            print(err)
            alert(str(err))

        self.timer_camera.stop()
        self.ui.stopButton.setEnabled(False)
        self.ui.rightDataButton.setEnabled(True)
        self.ui.leftDataButton.setEnabled(True)

        if len(self.left_temp_depth_list) == self.data_count and len(self.left_temp_rgb_list) == self.data_count \
                and len(self.right_temp_rgb_list) == self.data_count and len(
            self.right_temp_depth_list) == self.data_count:
            self.ui.exportButton.setEnabled(True)
            # alert(msg="全部采集完成，可以导出数据。",title="消息" , icon=QtWidgets.QMessageBox.Information)
            self.savePersonInfo()
        else:
            self.ui.exportButton.setEnabled(False)
